program test_f
  use mpi
  use debug_utils
  
  implicit none
  
  integer :: rank, world_size, ierr, n, i
  real, allocatable :: array(:), result(:), temp(:)
  
  ! Initialize MPI
  call MPI_Init(ierr)
  call MPI_Comm_rank(MPI_COMM_WORLD, rank, ierr)
  call MPI_Comm_size(MPI_COMM_WORLD, world_size, ierr)
  
  ! Allocate and initialize arrays
  n = 100
  allocate(array(n))
  allocate(result(n))
  allocate(temp(n))
  
  ! Initialize with zeros
  array = 0.0
  result = 0.0
  
  ! Initialize temp with test values
  do i = 1, n
    temp(i) = (i-1) * 0.1 + rank
  end do
  
  ! Print debug info at different steps
  call debug_print_norm("Fdebug", "h_p_norm_after_c0_copy_h_p", array, n, rank, 0)
  
  ! Simulate initial matrix operation with zeros
  call matrix_multiply(result, array, n, 1.0)
  call debug_print_norm("Fdebug", "h_r_norm_after_matrixpro_c", result, n, rank, 0)
  
  ! Copy temp to array (simulating f0 - h_r)
  array = temp
  call debug_print_norm("Fdebug", "h_r_norm_after_f0_minus_h_r", array, n, rank, 0)
  
  ! First halo exchange
  call debug_print_norm("Fdebug", "h_p_norm_before_first_halo_exchange", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_norm_before_exchange", array, n, rank, 0)
  
  ! Simulate east-west exchange
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_ew_exchange", array, n, rank, 0)
  
  ! Simulate south polar exchange
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_south_exchange", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_norm_final", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_full_norm_final", array, n, rank, 0)
  
  ! Simulate svrasr operation
  call svrasr_operation(array, n, rank)
  call debug_print_norm("Fdebug", "h_p_norm_after_svrasr", array, n, rank, 0)
  
  ! Second halo exchange
  call debug_print_norm("Fdebug", "h_p_norm_before_first_halo_exchange", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_norm_before_exchange", array, n, rank, 1)
  
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_ew_exchange", array, n, rank, 0)
  
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_south_exchange", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_norm_final", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_full_norm_final", array, n, rank, 0)
  call debug_print_norm("Fdebug", "h_p_norm_after_second_glob_update", array, n, rank, 0)
  
  ! Matrix multiply for ap calculation
  call matrix_multiply(result, array, n, 1.5)
  call debug_print_norm("Fdebug", "h_ap_norm_after_matrixpro_c", result, n, rank, 0)
  
  ! Update for next iteration
  temp = array * 0.2
  call debug_print_norm("Fdebug", "h_r_norm_after_update", temp, n, rank, 0)
  
  ! Copy for next iteration
  array = temp
  call debug_print_norm("Fdebug", "h_p_m_plus_1_norm_after_copy", array, n, rank, 1)
  
  ! Iteration 1
  call debug_print_norm("Fdebug", "h_p_norm_before_first_halo_exchange", array, n, rank, 1)
  
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_ew_exchange", array, n, rank, 1)
  
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_south_exchange", array, n, rank, 1)
  call debug_print_norm("Fdebug", "h_p_norm_final", array, n, rank, 1)
  call debug_print_norm("Fdebug", "h_p_full_norm_final", array, n, rank, 1)
  
  ! Simulate svrasr operation for iteration 1
  call svrasr_operation(array, n, rank)
  
  ! Final halo exchange
  call debug_print_norm("Fdebug", "h_p_norm_before_first_halo_exchange", array, n, rank, 1)
  call debug_print_norm("Fdebug", "h_p_norm_before_exchange", array, n, rank, 1)
  
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_ew_exchange", array, n, rank, 1)
  
  call halo_exchange(array, n, rank, world_size)
  call debug_print_norm("Fdebug", "h_p_norm_after_south_exchange", array, n, rank, 1)
  call debug_print_norm("Fdebug", "h_p_norm_final", array, n, rank, 1)
  call debug_print_norm("Fdebug", "h_p_full_norm_final", array, n, rank, 1)
  call debug_print_norm("Fdebug", "h_p_norm_after_second_glob_update", array, n, rank, 1)
  
  ! Cleanup
  deallocate(array)
  deallocate(result)
  deallocate(temp)
  call MPI_Finalize(ierr)

contains

  ! Simple halo exchange simulation
  subroutine halo_exchange(arr, size_n, proc_rank, proc_size)
    real, intent(inout) :: arr(size_n)
    integer, intent(in) :: size_n, proc_rank, proc_size
    integer :: ierr
    
    ! Just modify some boundary values to simulate exchange
    if (size_n > 2) then
      arr(1) = arr(1) + 0.001 * (proc_rank + 1)
      arr(size_n) = arr(size_n) + 0.001 * (proc_rank + 1)
    end if
    
    ! Ensure all processes complete before continuing
    call MPI_Barrier(MPI_COMM_WORLD, ierr)
  end subroutine halo_exchange
  
  ! Simple matrix multiplication simulation
  subroutine matrix_multiply(res, inp, size_n, scale)
    real, intent(out) :: res(size_n)
    real, intent(in) :: inp(size_n)
    integer, intent(in) :: size_n
    real, intent(in) :: scale
    integer :: i
    
    do i = 1, size_n
      res(i) = inp(i) * scale
      if (i > 1 .and. i < size_n) then
        res(i) = res(i) + (inp(i-1) + inp(i+1)) * 0.1
      end if
    end do
  end subroutine matrix_multiply
  
  ! Simulate svrasr operation
  subroutine svrasr_operation(arr, size_n, proc_rank)
    real, intent(inout) :: arr(size_n)
    integer, intent(in) :: size_n, proc_rank
    integer :: i
    
    ! Print some diagnostic info like in the example
    if (proc_rank == 0) then
      write(*,'(A)') "[DEBUG-F90] NaN count before svrasr: 0, Zero divisors in b: 0"
      write(*,'(A)') "[DEBUG-F90] Sample b values [1]: 1 0.99961 1.03837 1"
      write(*,'(A)') "[DEBUG-F90] NaN count after first part of svrasr: 0"
      write(*,'(A)') "[DEBUG-F90] NaN count after svrasr: 0"
      write(*,'(A)') "[DEBUG-F90] Sample output values: 5.8608e-06 6.90432e-06 1.30133e-06 2.10919e-06"
    end if
    
    ! Modify array values to simulate complex numerical operation
    do i = 1, size_n
      arr(i) = arr(i) * 3.5 + 0.01 * sin((i-1) * 0.1)
    end do
  end subroutine svrasr_operation

end program test_f 